﻿module BJK
(*
 * simplistic implementation of the new analysis for polynomial/linear bounds
 * over the "core language".
 *
 * To operate: evaluate Analyse(Cmd) where 
 * Cmd is a command, represented by an element of the datatype Cmd below.
 *
 *  See examples in examples.ml
 *)
 

(*
 * BASIC DEFINITIONS
 *)
open System

(* the same "D" is used for both analyses *)
module D = 
    type t = Dee of int   
    let Zero = Dee 0
    let One  = Dee 10
    let One' = Dee 15
    let Two  = Dee 20
    let Three = Dee 30
    let lub(Dee x, Dee y) = 
        Dee(Util.max(x,y))
        
    let sum(a,b) = 
        if (a=Three) || (b=Three) then Three
        elif (a=Zero) then b
        elif (b=Zero) then a
        else Two
    let prod(((Dee x) as a), ((Dee y) as b)) =
        if ((a=Zero) || (b = Zero)) then Zero else Dee(Util.max(x,y))
    let eql(Dee x, Dee y) = (x=y)
    let isOne(v) = (v = One || v = One')
    let toString((Dee x) as v) = 
        let digit = (x / 10).ToString()
        if v=One' then digit + "+" else digit 
                        
module Pattern =
    type Pat = | Cmd of int        (* command identifier *)
               | Seq of Pat * Pat  (* sequence of 2 commands *)
               | Loop of int * Pat (* Loop(command idnetifier, loop body pattern) *)
               | Pow of Pat * int  (* Pow(p,i) = p^i, i.e., i repetitions of the pattern p *)
               | Skip              (* Skip command gets its own pattern since it has special
                                      handling in the pattern product *)

    let rec eq(p1 : Pat, p2 : Pat) = match (p1,p2) with
                                     | (Cmd(i1),Cmd(i2)) -> i1 = i2
                                     | (Seq(p1,p2),Seq(p1',p2')) -> eq(p1,p1') && eq(p2,p2')
                                     | (Loop(i1,p1),Loop(i2,p2)) -> i1 = i2 && eq(p1,p2)
                                     | (Pow(p1,i1),Pow(p2,i2)) -> i1 = i2 && eq(p1,p2)
                                     | (Skip,Skip) -> true
                                     | _ -> false

    let naiveProd(p1 : Pat, p2 : Pat) = match (p1,p2) with
     (* skip is ignored *)              | (Skip,_) -> p2 
     (* skip is ignored *)              | (_,Skip) -> p1 
     (* Remove repetitions of loop body
        after loop correction *)        | (Loop(i,p),p') -> if eq(p,p') then Loop(i,p) else Seq(p1,p2)
     (* Concat two power patterns *)    | (Pow(p1',i1),Pow(p2',i2)) -> if eq(p1',p2') then Pow(p1',i1+i2) else Seq(p1,p2)
     (* increase power by 1 *)          | (Pow(p1',i1),p2') -> if eq(p1',p2') then Pow(p1',i1+1) else Seq(p1,p2)
     (* increase power by 1 *)          | (p2',Pow(p1',i1)) -> if eq(p1',p2') then Pow(p1',i1+1) else Seq(p1,p2)
                                        | _ -> if eq(p1,p2) then Pow(p1,2) else Seq(p1,p2)

    let rec reduce(p : Pat) = match p with
//                                | Seq((Loop(i,p) as pl),p') -> if (eq(p,p') then Loop(i,reduce(p)) 
//                                                               else naiveProd(reduce(pl),reduce(p'))
                                | Loop(i,p) -> Loop(i,reduce(p))
                                | Seq(p1,p2) -> naiveProd(reduce(p1),reduce(p2))
                                | _ -> p


    let prod(p1 : Pat, p2 : Pat) = 
        let p = naiveProd(p1,p2)
        let rec lfp f x = 
                      let fx = (f x)
                      if (eq(fx,x)) then x else lfp f fx
        lfp reduce p

    let rec size(p : Pat) = 1 + match p with 
                                | Seq(p1,p2) -> size(p1) + size(p2) (* sequence of 2 commands *)
                                | Loop(_,p') -> size(p') (* Loop(command idnetifier, loop body pattern) *)
                                | Pow(p',_) -> size(p')  (* Pow(p,i) = p^i, i.e., i repetitions of the pattern p *)
                                | _ -> 0

    let isSmaller(p1,p2) = size(p1) < size(p2)

    let smaller(p1,p2) = if isSmaller(p1,p2) then p1 else p2


(************************************************************)
(* sparse matrix implementation (or maybe it's a graph?) *)
module mat = 
    type DataFlow = (int * int * D.t * Pattern.Pat)
    type matrix = M of (DataFlow) list
         (* meaning: a SORTED list of entries (i,j,value)
           the sorting order is (j,i) to keep columns together *)

    let lstPrt(M m) = List.map (fun (i,j,v,p) -> (i,j,D.toString(v),p)) m;
 
    let prt(m) =          
        let rec prtInternal(entryList) = 
            match entryList with
            | [] -> ""
            | ((i, j, d,_)::t) -> (i).ToString() + "-(" + d + ")->" + (j).ToString() + ", " + prtInternal(t)
        prtInternal(lstPrt(m))

    let I(n:int,pat) = M(List.init n (fun i -> (i+1,i+1,D.One,pat)));

    let rec listsame = function
        | ([],[]) -> true
        | (_,[])  -> false
        | ([],_)  -> false
        | ((i1,j1,v1,_)::t1 as m1, (i2,j2,v2,_)::t2 as m2) -> 
            (i1,j1)=(i2,j2) && D.eql(v1,v2) && listsame(t1,t2)

    let same(M m1, M m2) = listsame(m1,m2);


    (* listsqcup merges sorted lists *)
    let rec listsqcup = function
        | ([],m2) -> m2
        | (m1,[]) -> m1
        | (((i1,j1,v1,p1)::t1 as m1), ((i2,j2,v2,p2)::t2 as m2)) ->
            if (i1,j1)=(i2,j2) then
                let p = match (v1,v2) with
                        | (D.Dee x1, D.Dee x2) -> Pattern.reduce(if x1 > x2 then p1 
                                                                 else if x2 > x1 then p2 
                                                                 else Pattern.smaller(p1,p2))

                (i1,j1,D.lub(v1,v2),p) :: listsqcup(t1,t2)
            elif j1<j2 || (j1=j2 && i1<i2) then (i1,j1,v1,p1) :: listsqcup(t1,m2)
            else (i2,j2,v2,p2) :: listsqcup(m1,t2)

    let rec sortLst = function
        | []              -> []
        | [ent]           -> [ent]
        | (e1::(e2::rst)) -> listsqcup([e1], sortLst(e2::rst)) 

    let sqcup(M m1, M m2) = M(listsqcup(sortLst(m1),sortLst(m2)));

    let prod(M m1, M m2) =
        M(List.fold
            (fun l1 l2 -> listsqcup(l1,l2))
            []
            (List.map
                (fun ((i1,j1,v,p1)) ->  (* multiply one entry from m1 by whole of m2*)
                    List.fold (
                        fun acc (i2,j2,u,p2) -> 
                            if j1=i2 then (i1,j2,D.prod(v,u),Pattern.prod(p1,p2))::acc else acc)
                        []
                        m2
                    )
                m1 
            )
        )
             

    (* lst2mat sorts its input - inefficiently *)
    let rec lst2mat = function
        | [] -> M([])
        | [ent] -> M([ent])
        | (e1::(e2::rst)) -> sqcup(M [e1], lst2mat(e2::rst))


    (* compute a matrix which is like (I n) except that column j is replaced
       with the provided 'col' *)
    let ins_in_I(n,j,col) =
        let ilist = match I(n,Pattern.Cmd(-1)) with
                    | M(l) -> l
                    
        M(List.fold
            (fun acc (i,jj,v,p) -> (( if jj=j then col else [(i,jj,v,p)] ) @ acc))
             []
             ilist     
         )
 
    let LC(l:int, nonlinear:D.t,cmdIdx : int) (M m) =
          M(listsqcup(m, 
                      List.choose
                        ( fun (i,j,v,p) -> 
                            if j=l && i<>l then raise (Util.Fail("BJK.mat.LC: loop variable X" + j.ToString() + " modified inside loop body."))
                            elif i=j then
                                if D.eql(v,D.One') then Some (l,j,nonlinear,Pattern.Loop(cmdIdx,p)) 
                                elif D.eql(v,D.Two) then Some (l,j,D.Three,Pattern.Loop(cmdIdx,p))
                                else None
                            else None)
                        m) )

    let ones(M(m)) = List.filter (fun (_,_,d,_) -> D.isOne(d)) m;
    (* the r(M) function: extract all pairs of One/One' entries *)

    let r(m) =
        let ones' = ones(m)
        List.fold  
            (fun acc (i,j,d,p1) -> 
                acc @ (List.choose
                            (fun (i',j',d,p2) -> if (i'<i || (i'=i && j'<j)) then Some((i,j),(i',j'),p1) else None) ones'
                      )
            )
            [] 
            ones'
 
    (* classify the variables into "properly bounded" (called SAFE in the paper)
       and the rest.  The Unbounded ones are given first!! *)
    let classify(M m, vars) :(int list)*(int list) = 
         List.partition (fun i -> 
                   List.exists (fun (_,j,d,_) -> (j=i && D.eql(d,D.Three))) m)
                        ([1..vars])

(* module mat *)
(************************************************************)


(* DFR - a module for the SOP method.  Using the "sparse matrix"
 * representation.
 *)

module sop_bar = 
    type PosPar = (int * int) * (int * int) * Pattern.Pat
    
    type DFR = Dfr of mat.matrix * (PosPar list)
     
    (*  the identity DFR *)
    let dfrI(n) =
        let idMat = mat.I(n)
        Dfr(idMat, mat.r(idMat))

    let prtR(Dfr(_,r)) = 
        let rec prtRList = function
            | ([])                   -> ""
            | ( ((i,j),(i',j'),_)::t ) -> "(" + i.ToString() + "->" + j.ToString() + ", " + (i').ToString() + "->" + (j').ToString() + "), " + prtRList(t);
        prtRList(r)

    let prtDfr((Dfr(M,R)) as d) = "\n    M=[ " + mat.prt(M) + "]\n    R=[ " + prtR(d) + "]";

    (* functions for sorting/merging R's (represented as sorted PosPar lists *) 
    let lexlt(((a,b),(c,d),_) : PosPar, ((aa,bb),(cc,dd),_)) =
      if (b<bb) then true 
      elif (b>bb) then false 
      elif (a<aa) then true
      elif (a>aa) then false 
      elif (d<dd) then true 
      elif (d>dd) then false 
      else (c<cc);

    let lexEq(((a,b),(c,d),_) : PosPar, ((aa,bb),(cc,dd),_)) =
      if (b<bb) then false 
      elif (b>bb) then false 
      elif (a<aa) then false
      elif (a>aa) then false 
      elif (d<dd) then false 
      elif (d>dd) then false 
      else (c=cc); 

    let rec merge = function
        | ([],ys) -> ys 
        | (xs,[]) -> xs
        | (x::xs, y::ys) -> if lexlt(x,y) then x::merge(xs,  y::ys)
                            elif lexEq(x,y) then merge(x::xs, ys) else y::merge(x::xs,  ys)
                  
    (* sort(already-sorted, yet-not-sorted) *)
    let rec sort = function
        | (xs, [])    -> xs
        | (xs, y::ys) -> let rec insert = function
                             | [] -> [y] 
                             | (x::xs) -> if lexlt(x,y) then x::insert(xs)
                                          elif lexEq(x,y) then x::xs else y::x::xs
                         sort(insert(xs), ys)

    let rec filterRDuplicates = function
        | [] -> []
        | [r] -> [r]
        | ((((a,b),(c,d),p) as r1)::(((a',b'),(c',d'),p') as r2)::rs) -> 
                if (a=a' && b=b' && c=c' && d=d') then ((a,b),(c,d),Pattern.smaller(p,p'))::filterRDuplicates(rs) 
                                                  else r1::filterRDuplicates(r2::rs)
         
    (* remove from r those entries which do not correspond to a 1/1+ *)
    let filterR(r: PosPar list, m: mat.matrix) =
        let ones' = mat.ones(m)
        let filtered = List.filter 
                             (fun ((a,b),(c,d),_) -> (List.exists (fun (i,j,_,_) -> (i,j)=(a,b)) ones') && (List.exists (fun (i,j,_,_) -> (i,j)=(c,d)) ones') )
                             r
        let order x y = if lexlt(x,y) then 1 else -1
        let sorted = List.sortWith order filtered
        filterRDuplicates(sorted)


                          
    (* operations on DFRs *)    
    let dfrEmpty = function
        | (Dfr(mat.M([]),[])) -> true
        | _                   -> false

               
    let dfrsame(Dfr(M1,R1),Dfr(M2,R2)) =
        let rec rSame = function
            | ([],[]) -> true
            | ((a,b,_)::t1,(a',b',_)::t2) -> a=a' && b=b' && rSame(t1,t2)
            | _ -> false

        mat.same(M1,M2) && rSame(R1,R2)
                          
    (* product of DFR's *)

    (* first, product without handling 'diamonds' *)
    let rec naiveProd = function
        | ([((a,c),(b,d),p1)],[((e,g),(f,h),p2)]) ->
                let p = Pattern.prod(p1,p2)
                sort ([], 
                      (if c=d && d=f then [((a,h),(b,h),p)] else [])  @
                      (if c=d && d=e then [((a,g),(b,g),p)] else [])  @
                      (if c=e && d=f then [((a,g),(b,h),p)] else [])  @
                      (if c=f && d=e then [((a,h),(b,g),p)] else [])  @
                      (if c=e && e=f then [((a,g),(a,h),p)] else [])  @
                      (if d=e && e=f then [((b,g),(b,h),p)] else []) )
        | (R1, R2) ->
                List.fold (fun Ls (p:PosPar) -> 
                    merge(List.fold (fun Ls (q:PosPar) -> 
                                        merge(naiveProd([p],[q]),Ls)) 
                                     [] 
                                     R2, 
                          Ls)) 
                          [] 
                          R1

    (* mine R: splits a result of naiveprod into two lists, one containing the positions
     that have to become 2 because of diamonds, the other is the rest of the list,
     to become the R of the product *)
    let rec mine = function
        | ([]: PosPar list)           -> ([],[])
        | ((((a,c),(b,d),p) as h)::t) ->
                let (diamonds, rest) = mine(t) 
                if (a,c)=(b,d) then ( (a,c,D.Two,p)::diamonds, rest )
                else ( diamonds, h::rest )
             

    (* another part of product computation:   
       parts of R' created by a 1 in the lhs and a fanout in the rhs. *)    
    let prod1(lhsM:mat.matrix, rhsR:PosPar list) =
        let ones' = mat.ones(lhsM)
        List.fold 
            (fun acc ((s,j'),(t,j''),p1) -> 
                    merge (acc,
                           (List.choose (fun (i,j,_,p2) -> 
                                            if (s=t) && (s=j)
                                            then Some ((i,j'),(i,j''),Pattern.prod(p2,p1))
                                            else None)
                                         ones')))
            [] rhsR
           
    (* the symmetric thing:
       parts of R' created by a 1 in the rhs and a fanin in the rhs. *)    
    let prod2(lhsR:PosPar list, rhsM:mat.matrix) =
        let ones' = mat.ones(rhsM)
        List.fold (fun acc (s,t,_,p1) -> 
                        merge (acc,
                                (List.choose (fun ((i,j),(i',j'),p2) -> 
                                                    if (j=j') && (j=s)
                                                    then Some ((i,t),(i',t),Pattern.prod(p2,p1))
                                                    else None)
                                              lhsR)))
                   [] ones'

      
    let prod(Dfr(A,Ra), Dfr(B,Rb)) =
      let p1 = prod1(A,Rb)
      let p2 = prod2(Ra,B)
      let (d,r) = mine(naiveProd(Ra,Rb));
      let m' = mat.sqcup(mat.prod(A,B), mat.lst2mat(d))
      let m'' = merge(r,p1)
      let m''' = merge(m'',p2)
      let filtered = filterR(m''',m')
      Dfr (m' , filtered)

      
    (* join of DFRs *)

    let sqcup(Dfr(M1,R1), Dfr(M2,R2)) =
        let M = mat.sqcup(M1,M2)
        Dfr(M, filterR(merge(R1,R2), M))

     

    
     (*******NB this is without the iteration count ************)
    let loop(l,vars,nonlinear,MR,cmdIdx) = 
        (* "loop closure" *)
        let rec fix f x =
            let x' = f x
            if dfrsame(x,x') then x else fix f x'

        let loopAux LC (Dfr(M,R)) = 
            let tmpM = LC(M)
            let tmpDFR = Dfr(tmpM, filterR(R,tmpM))
            prod(tmpDFR,tmpDFR)

        fix (loopAux (mat.LC (l,nonlinear,cmdIdx))) (sqcup(dfrI(vars,Pattern.Skip),MR))

 (* module sop_bar *)


(*
 * THE ANALYSIS ALGORITHM
 *)
(* ANALYSIS OF EXPRESSIONS - in two versions *)
module Analysis =

    (* the purpose of these functions is to abstract an expression into a 
       single column matrix, sorted! *)

    (*Input to analysis type *)
    type ANALYSIS_PARMS = 
        abstract member nonlinear_Dvalue: D.t
        abstract member name: string


    type analysis(analysis_kind : ANALYSIS_PARMS) = 
         (* abstract addition of 'columns' *)
        let add' d = D.lub(D.One',d);
        let add'' (i,d) = (i, add' d);

        let rec add = function
            | ([], ys) -> List.map add'' ys
            | (xs, []) -> List.map add'' xs
            | ((((i,d)::xt)as xs), (((i',d')::yt) as ys)) -> 
                    if i=i' then (i,D.sum(d,d')) :: add(xt,yt)
                    elif i<i' then (i,add' d) :: add(xt,ys)
                    else (i',add' d') :: add(xs,yt)


        (* utility functions: collect the variables in an Expr, or find the highest *)
        let rec merge = function
            | ([],ys) -> ys
            | (xs,[]) -> xs
            | (((x::xt) as xs), ((y::yt) as ys)) ->
                    if x=y then x::(merge(xt,yt))
                    elif x<y then x::(merge(xt,ys))
                    else y::(merge(xs,yt)) ;
            
        let rec highest_var_inE = function
            | Core.Var i         -> i
            | Core.Plus(e1, e2)  -> Util.max(highest_var_inE e1, highest_var_inE e2)
            | Core.Times(e1, e2) -> Util.max(highest_var_inE e1, highest_var_inE e2)
            | Core.Zero          -> 1;
      
        let rec vars = function
            | Core.Var i         -> [i]
            | Core.Plus(e1, e2)  -> merge(vars e1, vars e2)
            | Core.Times(e1, e2) -> merge(vars e1, vars e2)
            | Core.Zero          -> []

        let rec highest_var = function 
            | Core.Skip          -> 1
            | Core.Asgn(j,e)     -> Util.max(j, highest_var_inE e)
            | Core.Seq(C1,C2)    -> Util.max(highest_var C1, highest_var C2) 
            | Core.Choice(C1,C2) -> Util.max(highest_var C1, highest_var C2)
            | Core.Loop(l,C)     -> Util.max(l, highest_var C);

     
     
        (* analyse an expression.  *)
        let rec analyseE = function
            | Core.Var i                 -> [(i,D.One)]
            | Core.Plus(e1, e2)          -> add(analyseE e1, analyseE e2)
            | ((Core.Times( _, _)) as e) -> List.map (fun i -> (i,analysis_kind.nonlinear_Dvalue)) (vars e)
            | Core.Zero                  -> Util.fail("BJK.Analysis.analysis: Zero assignment is unsupported.")


        (* main function: analyse a command *)
        member x.analyse(c) = 
            let vars = highest_var c;  

            let rec analyseC = function
                | Core.Skip          -> sop_bar.dfrI(vars,Pattern.Cmd(-1))
                | Core.Asgn(j,e)     ->
                        let col = (List.map (fun (i,r) -> (i,j,r,Pattern.Cmd(-1))) (analyseE e));
                        let m = mat.ins_in_I(vars,j,col)
                        sop_bar.Dfr(m,mat.r(m))
                | Core.Seq(C1,C2)    -> sop_bar.prod(analyseC C1, analyseC C2)
                | Core.Choice(C1,C2) -> sop_bar.sqcup(analyseC C1, analyseC C2)
                | Core.Loop(l,C)     -> sop_bar.loop(l,vars,analysis_kind.nonlinear_Dvalue,(analyseC C),-1);
            
            analyseC c
                                                        
        member x.resultText(c, dfr) =   
            let vars = highest_var c 
            match dfr with
                sop_bar.Dfr(M, R) ->
                    let rec prt_ints = function
                        | [] -> "\n"
                        | i::rest -> (i.ToString()) + (if rest<>[] then ", " + (prt_ints rest) else "\n")
    
                    let (nonbounded, bounded) = mat.classify(M,vars)
                    (sop_bar.prtDfr(dfr) + "\n\nVARS: " + (vars.ToString()) + "\n" + analysis_kind.name + ": " +
                     (prt_ints bounded) + ("NON-" + analysis_kind.name + ": ") + (prt_ints nonbounded)) 

    (* analysis *)

    type PolyParams() =
         interface ANALYSIS_PARMS with
             member this.nonlinear_Dvalue = D.Two
             member this.name = "POLYNOMIAL"

    type LinParams() =
         interface ANALYSIS_PARMS with
             member this.nonlinear_Dvalue = D.Three
             member this.name = "LINEAR"
             
    let Pol = analysis(PolyParams())

    let Lin = analysis(LinParams());


